import rospy
from geometry_msgs.msg import PoseStamped, Vector3Stamped
from nav_msgs.msg import Odometry
import sys
import numpy as np
from gazebo_msgs.msg import ModelStates
from pyquaternion import Quaternion

vehicle_type = sys.argv[1]
vehicle_num = int(sys.argv[2])
multi_pose_pub = [None]*vehicle_num
multi_cam_pose_pub = [None]*vehicle_num
multi_odom_pub = [None]*vehicle_num
multi_speed_pub = [None]*vehicle_num
multi_local_pose = [PoseStamped() for i in range(vehicle_num)]
multi_camera_pose = [PoseStamped() for i in range(vehicle_num)]
multi_local_odom = [Odometry() for i in range(vehicle_num)]
multi_speed = [Vector3Stamped() for i in range(vehicle_num)]
camera_extrinsic = Odometry()
laser_extrinsic = Odometry()

def add_extrinsic():
    rotate_matrix = np.array(3)
    quat = Quaternion(rotate_matrix)
    camera_extrinsic.header.stamp = rospy.Time().now()
    camera_extrinsic.pose.pose.position.x = -0.365
    camera_extrinsic.pose.pose.position.y = 0.0
    camera_extrinsic.pose.pose.position.z = -0.84
    camera_extrinsic.pose.pose.orientation = quat
    
    laser_extrinsic.header.stamp = rospy.Time().now()
    laser_extrinsic.pose.pose.position.x = -0.39
    laser_extrinsic.pose.pose.position.y = 0.0
    laser_extrinsic.pose.pose.position.z = -0.65
    laser_extrinsic.pose.pose.orientation = quat

def gazebo_model_state_callback(msg):
    for vehicle_id in range(vehicle_num):
        id = msg.name.index(vehicle_type+'_'+str(vehicle_id))
        multi_local_pose[vehicle_id].header.stamp = rospy.Time().now()
        multi_local_pose[vehicle_id].header.frame_id = 'world'
        multi_local_pose[vehicle_id].pose = msg.pose[id]
        
        multi_camera_pose[vehicle_id].header.stamp = rospy.Time().now()
        multi_camera_pose[vehicle_id].header.frame_id = 'world'
        multi_camera_pose[vehicle_id].pose.position.x = msg.pose[id].position.x + 0.365
        multi_camera_pose[vehicle_id].pose.position.y = msg.pose[id].position.y + 0.0
        multi_camera_pose[vehicle_id].pose.position.z = msg.pose[id].position.z + 0.84
        
        multi_local_odom[vehicle_id].header.stamp = rospy.Time().now()
        multi_local_odom[vehicle_id].header.frame_id = 'world'
        multi_local_odom[vehicle_id].pose.pose = msg.pose[id]
        multi_local_odom[vehicle_id].twist.twist = msg.twist[id]
        
        multi_speed[vehicle_id].header.stamp = rospy.Time().now()
        multi_speed[vehicle_id].header.frame_id = 'world'
        multi_speed[vehicle_id].vector = msg.twist[id]

if __name__ == '__main__':
    rospy.init_node('multi_'+vehicle_type+'_odom')
    gazebo_model_state_sub = rospy.Subscriber("/gazebo/model_states", ModelStates, gazebo_model_state_callback,queue_size=1)
    for i in range(vehicle_num):
        multi_pose_pub[i] = rospy.Publisher(vehicle_type+'_'+str(i)+'/vision_pose/pose', PoseStamped, queue_size=1)
        multi_odom_pub[i] = rospy.Publisher(vehicle_type+'_'+str(i)+'/visual_slam/odom', Odometry, queue_size=1)
        multi_speed_pub[i] = rospy.Publisher(vehicle_type+'_'+str(i)+'/vision_speed/speed', Vector3Stamped, queue_size=1)
        multi_cam_pose_pub[i] = rospy.Publisher(vehicle_type+'_'+str(i)+'/vision_pose/camera_pose', PoseStamped, queue_size=1)
        
        print("Get " + vehicle_type + "_" + str(i) + " groundtruth pose")
    camera_extrinsic_pub = rospy.Publisher("/vins_estimator/extrinsic", Odometry, queue_size=1)
    laser_extrinsic_pub = rospy.Publisher("/aloam_estimator/extrinsic", Odometry, queue_size=1)
    add_extrinsic()
    
    rate = rospy.Rate(30)
    s = set()
    while not rospy.is_shutdown():
        for i in range(vehicle_num):
            multi_pose_pub[i].publish(multi_local_pose[i])
            multi_odom_pub[i].publish(multi_local_odom[i])
            multi_speed_pub[i].publish(multi_speed[i])
            multi_cam_pose_pub[i].publish(multi_camera_pose[i])
            
            if i not in s:
                print(f'Publish PoseStamp GroundTruth Topic: {vehicle_type}_{i}/vision_pos/pos')
                print(f'Publish Odomtry GroundTruth Topic: {vehicle_type}_{i}/visual_slam/odom')
            s.add(i)
        camera_extrinsic_pub.publish(camera_extrinsic)
        laser_extrinsic_pub.publish(laser_extrinsic)
        try:
            rate.sleep()
        except:
            continue
